#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/5/24 10:33
@desc: 
"""
import os
from pathlib import Path

from utils.log import get_console_logger

__all__ = ['ScanFinder']


class ScanFinder(object):
    """
    基本原始图像数据读取
    """

    def __init__(self, root,
                 recurrence=True,
                 max_number=1e12,
                 display=100,
                 logger=None):
        """

        Parameters
        ----------
        root: 根目录
        recurrence: 是否递归搜索子目录
        max_number: 读入的最大样本数
        display: 显示间隔
        logger: 日志对象
        """
        if not Path(root).exists():
            raise NotADirectoryError

        # 使用栈保存目录
        self._directory_stack = [Path(root)]

        # 保存中间数据
        self._buf_data = None
        self._recurrence = recurrence

        self._max_number = max_number
        self._count = 0
        self._display = display

        if logger is not None:
            self._logger = logger
        else:
            self._logger = get_console_logger('FileReader')

    def __iter__(self):
        while 0 < len(self._directory_stack):
            path = self._directory_stack.pop()
            for file in path.glob('*'):
                if file.is_file():
                    if self._filter(file):
                        yield self._buf_data
                elif self._recurrence and file.is_dir():
                    self._directory_stack.append(file)
                else:
                    raise TypeError('不能确定的文件类型！')

    def _has_next(self):
        """
        判断是否已经超过最大读取数量

        Returns
        -------

        """
        if self._count % self._display == 0:
            self._logger.info(
                '正在读取第 {} 个文件 ...'.format(self._count))
        if self._count < self._max_number:
            self._count += 1
        else:
            raise StopIteration

    def _filter(self, file):
        """
        find the file which contains the LIDC-IDRI dicom scan

        Parameters
        ----------
        file

        Returns
        -------

        """
        if file.suffix == '.xml':
            parent = str(file.parent)
            if len(os.listdir(parent)) > 80:
                self._buf_data = {'dcm': parent,
                                  'xml': str(file)}
                return True
        else:
            self._buf_data = None
            return True
